from numba import jit
import numpy as np
from numba import cuda
import torch
@jit(nopython=True)
def square_func(x):
    return x * x

x = 112.0
print(square_func(x))


@jit('float64(float64)', nopython=True)  # 指定输入和输出类型，括号内的是参数类型，括号外的是返回值类型
def square_float_func(x):
    return x * x


x = 112.0
print(square_float_func(x))


@cuda.jit
def matrix_add(A, B, C, m, n):
    row, col = cuda.grid(2)
    if row < m and col < n:
        C[row, col] = A[row, col] + B[row, col]


m, n = 1024, 1024
A = np.random.rand(m, n).astype(np.float32)
B = np.random.rand(m, n).astype(np.float32)

C = np.zeros_like(A)  # 创建与A形状相同的0数组
threads_per_block = (16, 16)
blocks_per_grid = (m // threads_per_block[0], n // threads_per_block[1])
matrix_add[blocks_per_grid, threads_per_block](A, B, C, m, n)
print(C)

print(torch.cuda.is_available())